#pragma once

#include "RungeKutta.h"
#include "Projection.h"

/**
 * @file INSE.h
 * @author xingyifan
 * @date 2022-08-31 22:58
 *
 * @description: 不可压 NS 方程求解器
 */

template <int Dim>
class INSESolver
{
public:
    INSESolver(FuncXt p, Real k, Real nu) : t_(0), p_(p), k_(k), nu_(nu), msolver_("I-Laplacian", k * nu * 0.25) {}

    /**
     * @file INSE.h
     * @author xingyifan
     * @date 2022-08-31 22:58
     *
     * @description: 传入初始速度，空的存放压强的容器，推进次数
     */

    Real solve(GridX<Dim> &U, Grid<Dim> &P, int N)
    {
        // 开始计时
        clock_t start = clock();

        // 循环计算每一步，移动 N 次
        for (int i = 0; i < N; i++)
        {
            shift(U);
        }

        // 最后求解 p 的 Poisson 方程
        pressure(U, P);

        // 结束计时
        Real t = Real(clock() - start) / CLOCKS_PER_SEC;

        // 将时间归零
        t_ = 0;

        return t;
    }

private:
    /**
     * @file INSE.h
     * @author xingyifan
     * @date 2022-08-31 23:08
     *
     * @description: 通过速度网格推进
     */

    void shift(GridX<Dim> &U)
    {
        // 存放算子结果
        std::vector<GridX<Dim>> Div(6, GridX<Dim>(Dim));
        std::vector<GridX<Dim>> Lap(6, GridX<Dim>(Dim));
        std::vector<GridX<Dim>> Pro(6, GridX<Dim>(Dim));

        // 创建算子
        Laplacian<Dim> laplacian;
        Divergence<Dim> divergence;
        Projection<Dim> projection;

        // 填充器
        GhostFiller<Dim> &ghostfiller = msolver_.ghostfiller_;

        // RK 方法的系数
        Matrix &AI = imexrungekutta_.AI_;
        Matrix &AE = imexrungekutta_.AE_;
        Vector &b = imexrungekutta_.b_;
        Vector &c = imexrungekutta_.c_;

        // 临时容器
        GridX<Dim> Phi = U;

        // 循环求解 (I-k nu gamma L)X = F
        for (int i = 1; i < 6; i++)
        {
            // 先计算 D<uu>
            Div[i - 1] = std::move(divergence(Phi));

            // 然后投影
            Pro[i - 1] = std::move(projection(Div[i - 1]));

            // 在每个方向上计算
            for (int d = 0; d < Dim; d++)
            {
                // 记录 F
                Grid<Dim> F = U[d];
                for (int j = 0; j < i - 1; j++)
                {
                    F += TensorOp::sub(Lap[j][d] * (AI[i * 6 + j] * nu_ * k_), Pro[j][d] * (AE[i * 6 + j] * k_));
                }

                // 计算 L<u>
                Lap[i - 1][d] = std::move(laplacian(Phi[d], ghostfiller));

                F += TensorOp::sub(Lap[i - 1][d] * (AI[i * 6 + i - 1] * nu_ * k_),
                                   Pro[i - 1][d] * (AE[i * 6 + i - 1] * k_));

                // 使用多重网格求解
                msolver_.solve(F, Phi[d]);
            }
        }

        // 计算最后一项，此时不需要投影
        Div[5] = std::move(divergence(Phi));

        // 在每个方向上计算
        for (int d = 0; d < Dim; d++)
        {
            Lap[5][d] = std::move(laplacian(Phi[d], ghostfiller));
        }

        // 加权求和
        for (int i = 0; i < 6; i++)
        {
            // 在各方向加权求和
            for (int d = 0; d < Dim; d++)
            {
                U[d] += TensorOp::sub(Lap[i][d] * (b[i] * nu_ * k_), Div[i][d] * (b[i] * k_));
            }
        }

        // 最后投影
        U = std::move(projection(U));

        // 时间推进
        t_ += k_;
    }

    /**
     * @file INSE.h
     * @author xingyifan
     * @date 2022-08-31 20:41
     *
     * @description: 压强的 Poisson 方程
     */

    void pressure(GridX<Dim> &U, Grid<Dim> &P)
    {
        // 创建算子
        Laplacian<Dim> laplacian;
        Divergence<Dim> divergence;

        // 新创建多重网格法，用 Laplacian 作为系数矩阵
        MultigridSolver<Dim> solver;

        // 获得填充器
        GhostFiller<Dim> &ghostfiller = solver.ghostfiller_;

        // 计算 D<uu>
        GridX<Dim> Div = std::move(divergence(U));

        // 计算 L<u>
        GridX<Dim> Lap(Dim);

        // 各个方向上计算
        for (int d = 0; d < Dim; d++)
        {
            // 计算 nu L<u> - D<uu>
            Lap[d] = std::move(TensorOp::sub(TensorOp::multiple(laplacian(U[d], ghostfiller), nu_), Div[d]));
        }

        // 散度算子作用
        Grid<Dim> F = std::move(divergence(Lap, ghostfiller));

        // 求解 Poisson 方程
        solver.solve(F, P);

        // 创建单位区域
        Domain<Dim> domain;

        // 创建积分算子
        Integrator<3, Dim> I;

        // 包装函数
        FuncX p = [=](Vector x)
        {
            return p_(x, t_);
        };

        // 计算精确值的积分平均和计算解积分平均之差，补足差
        // 周期边界条件下 Laplacian 奇异，需要补一个特解
        P += (Real)I.CellAverage(domain, P, p) - P[0];
    }

private:
    FuncXt p_; // 压力项

    Real t_;  // 当前时间
    Real k_;  // 时间步长
    Real nu_; // 扩散常数

    IMEXRungeKutta imexrungekutta_; // 隐式-显式迭代
    MultigridSolver<Dim> msolver_;  // 多重网格求解器
};